# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: jingyuan

import os
import tempfile
from collections import OrderedDict
from concurrent.futures import ProcessPoolExecutor, as_completed
from datetime import datetime

from absl import app, flags
from jinja2 import Environment, FileSystemLoader

import pandas as pd

import coin.research.read_strat_info as rsi
from coin.base.mail_util import send_html_mail, send_mail_with_attachments

from .error_check_handlers import *
from .on_log import run_from_strat_log

FLAGS = flags.FLAGS


def render(context):
  """Renders the mail content"""
  loader = FileSystemLoader(os.environ['PYTHONPATH'].split(":"))
  env = Environment(loader=loader)
  env.globals.update(__builtins__.__dict__)
  template = env.get_template('coin2/support/og_log_checker/templates/email_template.html')
  return template.render(context)


def run(machine, strategy_name, market_type, ccy, start, end):
  handlers = [eval(klass) if "FillPositionMismatch" not in klass
                          else FillPositionMismatch.parse_futures_info(klass, market_type, ccy)
                          for klass in FLAGS.check.split(';')]

  run_from_strat_log(machine, strategy_name, start, handlers, end)
  return strategy_name, machine, handlers


def convert_ts(ts):
  if len(ts) == 19 and ts.isdigit():
    return int(ts)
  return pd.to_datetime(ts).timestamp() * 1e9


def main(argv):
  start_ts, end_ts = convert_ts(FLAGS.start_ts), convert_ts(FLAGS.end_ts)
  start, end = datetime.fromtimestamp(start_ts * 1e-9), datetime.fromtimestamp(end_ts * 1e-9)

  strategy_name = FLAGS.strategy_name
  # TODO(jingyuan): fix multi ccy strategy info
  strategy_infos = rsi.get_strategy_info(trading_date=start.date(), strategy_name=strategy_name)

  context = {}
  attachments = {}
  seen = set()
  futures = []
  total_errors = 0
  with ProcessPoolExecutor(max_workers=FLAGS.cpu) as executor:
    for strategy_info in strategy_infos:
      strategy_name = strategy_info['strategy_name']
      market_type = strategy_info['market_type']
      ccy = strategy_info['symbol'].split('-')[-1]
      if strategy_name not in seen:
        seen.add(strategy_name)
      else:
        continue
      machine = strategy_info['machine']
      futures.append(executor.submit(run, machine, strategy_name, market_type, ccy, start, end))

    for fut in as_completed(futures):
      strategy_name, machine, handlers = fut.result()
      for handler_name, handler in zip(FLAGS.check.split(';'), handlers):
        if handler.errors and len(df := handler.to_df()) > 0:
          key = (strategy_name, machine, handler_name, len(df))
          print(key)
          total_errors += len(df)
          context[key] = df.iloc[:FLAGS.present_limit]
          df = df.assign(strategy_name=strategy_name,
                         machine=machine)
          if handler_name not in attachments:
            attachments[handler_name] = df
          else:
            attachments[handler_name] = attachments[handler_name].append(df)

  # sort by number of errors
  context = OrderedDict(sorted(context.items(), key=lambda x: -x[0][-1]))

  html = render({'context': context})

  if flags.FLAGS.is_daily_job:
    assert flags.FLAGS.daily_out_dir
    date = start.strftime("%Y%m%d")
    dir_name = os.path.join(flags.FLAGS.daily_out_dir, flags.FLAGS.strategy_name, date)
    if not os.path.exists(dir_name):
      os.makedirs(dir_name)
    for name, df in attachments.items():
      name = "".join([ch for ch in name if ch.isalnum()])
      filename = f"{dir_name}/{name}.csv"
      df.to_csv(filename)
    msg = html + '\n' + os.path.join(flags.FLAGS.link_root, dir_name)
    send_html_mail(
        "Error Checker",
        "error_check@joomo.io",
        FLAGS.report_email,
        f"{start.date()} {FLAGS.strategy_name} {FLAGS.check} {total_errors} errors",
        msg)
  else:
    filenames = []
    with tempfile.TemporaryDirectory() as tempdir:
      tempfile.TemporaryFile(dir=tempdir, suffix='.csv')
      for name, df in attachments.items():
        name = "".join([ch for ch in name if ch.isalnum()])
        filename = f"{tempdir}/{name}.csv"
        df.to_csv(filename)
        filenames.append(filename)

      send_mail_with_attachments(
          "Error Checker",
          "error_check@joomo.io",
          FLAGS.report_email,
          f"{start.date()} {FLAGS.strategy_name} {FLAGS.check} {total_errors} errors",
          "",
          filenames,
          html_text=html)


if __name__ == "__main__":
  flags.DEFINE_integer('present_limit', 5, 'how many rows in mail body')
  flags.DEFINE_integer('cpu', 16, '')
  flags.DEFINE_string('check', 'CancelNumber();OrderLifetime(120)', 'stuffs to check')
  flags.DEFINE_string('end_ts', '20200805-23', 'ns or any format supported by pd.to_datetime')
  flags.DEFINE_string('report_email', '', '')
  flags.DEFINE_string('start_ts', '20200805-235959', '')
  flags.DEFINE_string('strategy_name', 'basis_smm_binance_perpetual_okex_quarter', '')
  flags.DEFINE_string('daily_out_dir', 'csv_out/daily', '')
  flags.DEFINE_string('link_root', 'http://jenkins-hft.corp.prestolabs.io/'
                      'view/coin/job/coin2_og_log_error_checker/ws/coin/', '')
  flags.DEFINE_boolean('is_daily_job', True, '')
  app.run(main)
